import os

import PIL.Image
import math
import numpy as np
import torch
from matplotlib import pyplot as plt
from torchvision.transforms import transforms

import silearn
import silearn.graph
from model.encoding_tree import Partitioning
from optimizer.enc.partitioning.propagation_others import OperatorPropagation
os.environ['CUDA_VISIBLE_DEVICES']  = "1"
img = PIL.Image.open("img_2.png")
img = np.array(img)[:,:,:3]
img = torch.tensor(img).double() / 255
target = "SE_M"


w, es, et = silearn.spatial_knn_graph(img.cuda(), 81, 9)
w = torch.exp2(-w / w.mean())
es,et =es.cuda(), et.cuda()
imgH = img.shape[0]
imgW = img.shape[1]
g = silearn.graph.GraphSparse(torch.cat([es.unsqueeze(1), et.unsqueeze(1)],dim = 1), w, n_vertices =imgH * imgW)
delta = torch.abs(es - et)
connectivity = (delta == 1) + (delta == imgW)
optim = OperatorPropagation(Partitioning(g, None), objective=target)
optim.perform(m_scale=0, min_com=100, adj_cover=connectivity)



# for i in range(12):
#     optim.perform(m_scale=2**i, min_com=10, re_compute=False)
# optim.perform(m_scale=4, min_com=2, re_compute=False)
# optim.perform(m_scale=16, min_com=2, re_compute=False)
# optim.perform(m_scale=64, min_com=2, re_compute=False)
# optim.perform(m_scale=256, min_com=2, re_compute=False)
# optim.perform(m_scale=1024, min_com=2, re_compute=False)

seg = optim.enc.node_id.cpu()
print(seg.max() + 1 )

torch.random.manual_seed(321)
imgc = torch.rand((seg.max() + 1, 3))
imgc2 = torch.rand((seg.max() + 1, 3))
imgs = imgc[seg].reshape(imgH, imgW, 3)
# imgs2 = img.permute(1, 2, 0).reshape(imgH, imgW, 3)
imgs = (img + imgs * 2) /3

imgcl = seg.reshape(1, 1, imgH, imgW).float()
kernel1 = torch.Tensor([-1, 1, 0]).reshape(1, 1, 1, 3)
kernel2 = torch.Tensor([-1, 1, 0]).reshape(1, 1, 3, 1)
import torch.nn.functional as F
imgcl = torch.max(
    torch.abs(F.conv2d(imgcl, kernel1, padding=(0, 1))) + torch.abs(F.conv2d(imgcl, kernel2, padding=(1, 0))), 1)

imgs[:, :, 0] += imgcl[0][0]
imgs[:, :, 1] -= imgcl[0][0]
imgs[:, :, 2] -= imgcl[0][0]

imgs = imgs.clamp(max=1, min=0)
imgs = np.array(imgs)
plt.imsave("output.png", imgs)
plt.axis("off")
plt.imshow(imgs)
plt.title(target)
# plt.imshow(img.numpy())
plt.show()
# graph =
